﻿#region Copyright information
// 
// Copyright © 2005-2013 Yongkee Cho. All rights reserved.
// 
// This code is a part of the Biological Object Library and governed under the terms of the
// GNU Lesser General  Public License (LGPL) version 2.1 which accompanies this distribution.
// For more information on the LGPL, please visit http://bol.codeplex.com/license.
// 
// - Filename: MarkovModelBase.cs
// - Author: Yongkee Cho
// - Email: yongkeecho@gmail.com
// - Date Created: 2013-01-24 4:34 PM
// - Last Modified: 2013-01-25 3:59 PM
// 
#endregion
using System;
using System.Collections.Generic;
using System.Linq;
using BOL.Linq;
using BOL.Linq.Probability;
using BOL.Maths.Distributions;

namespace BOL.Algorithms.StateMachines
{
    public abstract class MarkovModelBase<TState> : IMarkovChain<TState>
    {
        #region Public properties

        public IDictionary<TState, double> InitialDistribution { get; set; }
        public IDictionary<TState, IDistribution<TState>> TransitionDistribution { get; set; }

        #endregion

        #region Constructor

        protected MarkovModelBase(IDictionary<TState, double> initialDistribution, IDictionary<TState, IDistribution<TState>> transitionDistribution)
        {
            if (initialDistribution == null)
                throw new ArgumentNullException("initialDistribution");
            if (transitionDistribution == null)
                throw new ArgumentNullException("transitionDistribution");
            if (initialDistribution.Count == 0)
                throw new ArgumentOutOfRangeException("initialDistribution", "Initial distribution cannot be empty.");
            if (transitionDistribution.Count == 0)
                throw new ArgumentOutOfRangeException("transitionDistribution", "Transition distribution cannot be empty.");

            InitialDistribution = initialDistribution;
            TransitionDistribution = transitionDistribution;
        }

        #endregion

        #region Public methods

        public IEnumerable<TState> Generate(Random r, int length)
        {
            if (r == null)
                throw new ArgumentNullException("r");

            var current = InitialDistribution.Sample(r);
            for (var i = 0; i < length; ++i)
            {
                yield return current;

                if (!TransitionDistribution.ContainsKey(current))
                    yield break;

                current = TransitionDistribution[current].Sample(r);
            }
        }

        public void Train(IEnumerable<TState> source)
        {
            if (source == null)
                throw new ArgumentNullException("source");

            var array = source.ToArray();

            // Updates transition distribution
            var transitionDistribution = InitialDistribution.ToDictionary(x => x.Key, x => new List<TState>());
            foreach (var key in array.Window(2).Select(x => x.ToArray()))
                transitionDistribution[key[0]].Add(key[1]);
            TransitionDistribution.ForEach(x => x.Value.MaximumLikelihoodEstimate(transitionDistribution[x.Key]));
        }

        public void Train(IEnumerable<IEnumerable<TState>> source)
        {
            if (source == null)
                throw new ArgumentNullException("source");

            var list = source.Select(sequence => sequence.ToArray()).ToList();

            // Updates initial distribution
            var firsts = list.Select(x => x[0]).ToList();
            InitialDistribution = InitialDistribution.ToDictionary(x => x.Key, x => (double)firsts.Count(y => x.Key.Equals(y))).Normalize();

            // Updates transition distribution
            var transitionDistribution = InitialDistribution.ToDictionary(x => x.Key, x => new List<TState>());
            foreach (var key in list.SelectMany(sequence => sequence.Window(2).Select(window => window.ToArray())))
                transitionDistribution[key[0]].Add(key[1]);
            TransitionDistribution.ForEach(x => x.Value.MaximumLikelihoodEstimate(transitionDistribution[x.Key]));
        }

        #endregion

        #region IEquatable<IMarkovChain<TState>> implementation

        public bool Equals(IMarkovChain<TState> other)
        {
            if (other == null)
                throw new ArgumentNullException("other");

            return InitialDistribution.DictionaryEqual(other.InitialDistribution);
        }

        #endregion

        #region Object overriden

        public override int GetHashCode()
        {
            return InitialDistribution.OrderBy(x => x.Key).Aggregate(1, (hashCode, x) => hashCode ^ x.Key.GetHashCode());
        }

        public override bool Equals(object other)
        {
            if (other == null)
                throw new ArgumentNullException("other");

            if (!(other is IMarkovChain<TState>))
                throw new InvalidCastException("The 'other' argument is not a IMarkovChain<TState> object.");

            return Equals(other as IMarkovChain<TState>);
        }

        public override string ToString()
        {
            return String.Format("{0} {{States = {{{1}}}, Transitions = {{{2}}}}}",
                GetType().Name,
                String.Join(", ", InitialDistribution.Select(x => String.Format("{0} ({1:0.0000})", x.Key, x.Value))),
                String.Join(", ", TransitionDistribution.Select(x => String.Format("{0}->{1}", x.Key, x.Value))));
        }

        #endregion
    }
}